package com.dbflow5.processor;

import com.dbflow5.MapUtils;
import com.dbflow5.StringUtils;
import com.dbflow5.processor.definition.*;
import com.dbflow5.processor.definition.provider.ContentProviderDefinition;
import com.dbflow5.processor.definition.provider.TableEndpointDefinition;
import com.dbflow5.processor.utils.WriterUtils;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.TypeName;

import javax.annotation.processing.FilerException;
import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import javax.tools.Diagnostic;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Description: The main object graph during processing. This class collects all of the
 * com.dbflow5.processor classes and writes them to the corresponding database holders.
 */
public class ProcessorManager implements Handlers.Handler {

    public static ProcessorManager manager;

    public ProcessingEnvironment processingEnvironment;
    private final List<TypeName> uniqueDatabases = new ArrayList<>();
    private final Map<TypeName, TypeName> modelToDatabaseMap = new HashMap<>();
    public final LinkedHashMap<TypeName, TypeConverterDefinition> typeConverters = new LinkedHashMap<>();
    private final Map<TypeName, Map<Integer, List<MigrationDefinition>>> migrations = new HashMap<>();

    private final Map<TypeName, DatabaseObjectHolder> databaseDefinitionMap = new HashMap<>();
    private final Set<Handlers.AnnotatedHandler<?>> handlers = new HashSet<>();
    private final Map<TypeName, ContentProviderDefinition> providerMap = new HashMap<>();

    public ProcessorManager(ProcessingEnvironment processingEnvironment) {
        this.processingEnvironment = processingEnvironment;
        manager = this;

        messager = processingEnvironment.getMessager();
        typeUtils  = processingEnvironment.getTypeUtils();
        elements = processingEnvironment.getElementUtils();
    }

    public void addHandlers(Handlers.AnnotatedHandler<?>... containerHandlers) {
        handlers.addAll(Arrays.asList(containerHandlers));
    }

    public Messager messager;

    public Types typeUtils;

    public Elements elements;

    public void addDatabase(TypeName database) {
        if (!uniqueDatabases.contains(database)) {
            uniqueDatabases.add(database);
        }
    }

    public void addDatabaseDefinition(DatabaseDefinition databaseDefinition) {
        DatabaseObjectHolder holderDefinition = getOrPutDatabase(databaseDefinition.elementClassName);
        if(holderDefinition != null) {
            holderDefinition.databaseDefinition = databaseDefinition;
        }
    }

    public List<DatabaseObjectHolder> getDatabaseHolderDefinitionList() {
        return new ArrayList<>(databaseDefinitionMap.values());
    }

    public DatabaseObjectHolder getDatabaseHolderDefinition(TypeName databaseName) {
        return databaseDefinitionMap.get(databaseName);
    }

    public void addTypeConverterDefinition(TypeConverterDefinition definition) {
        typeConverters.put(definition.modelTypeName, definition);
    }

    public TypeConverterDefinition getTypeConverterDefinition(TypeName typeName) {
        return typeConverters.get(typeName);
    }

    public void addModelToDatabase(TypeName modelType, TypeName databaseName) {
        if(modelType != null) {
            addDatabase(databaseName);
            modelToDatabaseMap.put(modelType, databaseName);
        }
    }

    public void addQueryModelDefinition(QueryModelDefinition queryModelDefinition) {
        ClassName it = queryModelDefinition.elementClassName;
        if(it != null) {
            getOrPutDatabase(queryModelDefinition.associationalBehavior().databaseTypeName).queryModelDefinitionMap.put(it, queryModelDefinition);
        }
    }

    public void addTableDefinition(TableDefinition tableDefinition) {
        if(tableDefinition.elementClassName != null) {
            DatabaseObjectHolder holderDefinition = getOrPutDatabase(tableDefinition.associationalBehavior().databaseTypeName);
            if(holderDefinition.tableDefinitionMap != null) {
                holderDefinition.tableDefinitionMap.put(tableDefinition.elementClassName, tableDefinition);
            }
            if(holderDefinition.tableNameMap != null) {
                String tableName = tableDefinition.associationalBehavior().name;
                if (holderDefinition.tableNameMap.containsKey(tableName)) {
                    logError("Found duplicate table "+tableName+" " +
                            "for database " + holderDefinition.databaseDefinition.elementName);
                } else {
                    holderDefinition.tableNameMap.put(tableName, tableDefinition);
                }
            }
        }
    }

    public void addManyToManyDefinition(ManyToManyDefinition manyToManyDefinition) {
        DatabaseObjectHolder databaseHolderDefinition = getOrPutDatabase(manyToManyDefinition.databaseTypeName);
        if(databaseHolderDefinition.manyToManyDefinitionMap != null && manyToManyDefinition.elementClassName != null) {
            MapUtils.getOrPut(databaseHolderDefinition.manyToManyDefinitionMap, manyToManyDefinition.elementClassName, new ArrayList<>()).add(manyToManyDefinition);
        }
    }

    public TableDefinition getTableDefinition(TypeName databaseName, TypeName typeName) {
        return getOrPutDatabase(databaseName).tableDefinitionMap != null? getOrPutDatabase(databaseName).tableDefinitionMap.get(typeName): null;
    }

    public QueryModelDefinition getQueryModelDefinition(TypeName databaseName, TypeName typeName) {
        Map<TypeName, QueryModelDefinition> map = getOrPutDatabase(databaseName).queryModelDefinitionMap;
        if(map != null) {
            return map.get(typeName);
        }
        return null;
    }

    public ModelViewDefinition getModelViewDefinition(TypeName databaseName, TypeName typeName) {
        Map<TypeName, ModelViewDefinition> map = getOrPutDatabase(databaseName).modelViewDefinitionMap;
        return map != null? map.get(typeName) : null;
    }

    public EntityDefinition getReferenceDefinition(TypeName databaseName, TypeName typeName) {
        TableDefinition tableDefinition = getTableDefinition(databaseName, typeName);
        if(tableDefinition != null){
            return tableDefinition;
        }else {
            QueryModelDefinition queryModelDefinition = getQueryModelDefinition(databaseName, typeName);
            if(queryModelDefinition != null) {
                return queryModelDefinition;
            }else {
                return getModelViewDefinition(databaseName, typeName);
            }
        }
    }

    public void addModelViewDefinition(ModelViewDefinition modelViewDefinition) {
        if(modelViewDefinition.elementClassName != null) {
            DatabaseObjectHolder holder = getOrPutDatabase(modelViewDefinition.associationalBehavior().databaseTypeName);
            if(holder != null && holder.modelViewDefinitionMap != null) {
                holder.modelViewDefinitionMap.put(modelViewDefinition.elementClassName, modelViewDefinition);
            }

        }
    }

    public List<TypeConverterDefinition> getTypeConverters() {
        List<TypeConverterDefinition> list = new ArrayList<>(typeConverters.values());
        list.sort(Comparator.comparing(o -> o.modelTypeName.toString()));
        return list;
    }

    public List<TableDefinition> getTableDefinitions(TypeName databaseName) {
        List<TableDefinition> list = new ArrayList<>();
        DatabaseObjectHolder databaseHolderDefinition = getOrPutDatabase(databaseName);
        if(databaseHolderDefinition.tableDefinitionMap != null) {
            list = new ArrayList<>(databaseHolderDefinition.tableDefinitionMap.values());
            list.sort(Comparator.comparing(tableDefinition -> tableDefinition.outputClassName.simpleName()));
        }

        return list;
    }

    public void setTableDefinitions(Map<TypeName, TableDefinition> tableDefinitionSet, TypeName databaseName) {
        DatabaseObjectHolder databaseDefinition = getOrPutDatabase(databaseName);
        databaseDefinition.tableDefinitionMap = tableDefinitionSet;
    }

    public List<ModelViewDefinition> getModelViewDefinitions(TypeName databaseName) {
        List<ModelViewDefinition> list = new ArrayList<>();
        DatabaseObjectHolder databaseDefinition = getOrPutDatabase(databaseName);
        if(databaseDefinition.modelViewDefinitionMap != null){
            list = new ArrayList<>(databaseDefinition.modelViewDefinitionMap.values());
            list.sort(Comparator.comparing(o -> o.outputClassName.simpleName()));
            list.sort(Comparator.comparing(o -> o.priority));
        }

        return list;
    }

    public void setModelViewDefinitions(Map<TypeName, ModelViewDefinition> modelViewDefinitionMap, ClassName elementClassName) {
        DatabaseObjectHolder databaseDefinition = getOrPutDatabase(elementClassName);
        databaseDefinition.modelViewDefinitionMap = modelViewDefinitionMap;
    }

    public List<QueryModelDefinition> getQueryModelDefinitions(TypeName databaseName) {
        List<QueryModelDefinition> list = new ArrayList<>();
        DatabaseObjectHolder databaseDefinition = getOrPutDatabase(databaseName);
        if(databaseDefinition.queryModelDefinitionMap != null) {
            list = new ArrayList<>(databaseDefinition.queryModelDefinitionMap.values());
            list.sort(Comparator.comparing(queryModelDefinition -> queryModelDefinition.outputClassName.simpleName()));
        }

        return list;
    }

    public void addMigrationDefinition(MigrationDefinition migrationDefinition) {
        Map<Integer, List<MigrationDefinition>> migrationDefinitionMap = MapUtils.getOrPut(migrations, migrationDefinition.databaseName, new HashMap<>());
        List<MigrationDefinition> migrationDefinitions = MapUtils.getOrPut(migrationDefinitionMap, migrationDefinition.version, new ArrayList<>());
        if (!migrationDefinitions.contains(migrationDefinition)) {
            migrationDefinitions.add(migrationDefinition);
        }
    }

    public Map<Integer, List<MigrationDefinition>> getMigrationsForDatabase(TypeName databaseName) {
        Map<Integer, List<MigrationDefinition>> map = migrations.get(databaseName);
        if(map != null) {
            return map;
        }else {
            return new HashMap<>();
        }
    }

    public void addContentProviderDefinition(ContentProviderDefinition contentProviderDefinition) {
        if(contentProviderDefinition.elementTypeName != null) {
            DatabaseObjectHolder holderDefinition = getOrPutDatabase(contentProviderDefinition.databaseTypeName);
            if(holderDefinition.providerMap != null) {
                holderDefinition.providerMap.put(contentProviderDefinition.elementTypeName, contentProviderDefinition);
            }
            providerMap.put(contentProviderDefinition.elementTypeName, contentProviderDefinition);
        }
    }

    public void putTableEndpointForProvider(TableEndpointDefinition tableEndpointDefinition) {
        ContentProviderDefinition contentProviderDefinition = providerMap.get(tableEndpointDefinition.contentProviderName);
        if (contentProviderDefinition == null) {
            logError("Content Provider "+tableEndpointDefinition.contentProviderName+" was not found for the @TableEndpoint " + tableEndpointDefinition.elementClassName);
        } else {
            contentProviderDefinition.endpointDefinitions.add(tableEndpointDefinition);
        }
    }

    public void logError(Class<?> callingClass, String error, Object... args) {
        messager.printMessage(Diagnostic.Kind.ERROR,
            String.format((callingClass!=null?callingClass.toString() : "")
                    // don't print this in logs.
                    .replace("(Java reflection is not available)", "")
             + ":" + (error != null? error.trim() : ""), args));
    }

    public void logError(String error) {
        logError( null, error);
    }

    public void logWarning(String error) {
        messager.printMessage(Diagnostic.Kind.WARNING, error != null? error : "");
    }

    public void logWarning(Class<?> callingClass, String error) {
        logWarning(callingClass + " : " + error);
    }

    private DatabaseObjectHolder getOrPutDatabase(TypeName databaseName) {
        return MapUtils.getOrPut(databaseDefinitionMap, databaseName, new DatabaseObjectHolder());
    }

    @Override
    public void handle(ProcessorManager processorManager, RoundEnvironment roundEnvironment) {
        handlers.forEach(annotatedHandler -> {
            annotatedHandler.handle(processorManager, roundEnvironment);
        });

        List<DatabaseObjectHolder> databaseHolderDefinitionList = getDatabaseHolderDefinitionList();
        databaseHolderDefinitionList.sort(Comparator.comparing(databaseObjectHolder -> databaseObjectHolder.databaseDefinition.outputClassName.simpleName()));
        for (DatabaseObjectHolder databaseHolderDefinition : databaseHolderDefinitionList) {
            try {

                if (databaseHolderDefinition.databaseDefinition == null) {
                    manager.logError(StringUtils.joinToString(databaseHolderDefinition.getMissingDBRefs(),"\n"));
                    continue;
                }

                Collection<List<ManyToManyDefinition>> manyToManyDefinitions = databaseHolderDefinition.manyToManyDefinitionMap.values();

                List<ManyToManyDefinition> sumList = new ArrayList<>();
                for (List<ManyToManyDefinition> list : manyToManyDefinitions) {
                    sumList.addAll(list);
                }
                List<ManyToManyDefinition> sortedList = sumList.stream().sorted(Comparator.comparing(manyToManyDefinition -> manyToManyDefinition.outputClassName.simpleName())).collect(Collectors.toList());
                for (ManyToManyDefinition manyToManyList : sortedList) {
                    manyToManyList.prepareForWrite();
                    WriterUtils.writeBaseDefinition(manyToManyList, processorManager);
                }

                // process all in next round.
                if (!manyToManyDefinitions.isEmpty()) {
                    manyToManyDefinitions.clear();
                    continue;
                }

                if (roundEnvironment.processingOver()) {
                    Validators.ContentProviderValidator validator = new Validators.ContentProviderValidator();
                    List<ContentProviderDefinition> contentProviderDefinitions = databaseHolderDefinition.providerMap.values()
                        .stream().sorted(Comparator.comparing(it -> it.outputClassName.simpleName())).collect(Collectors.toList());
                    contentProviderDefinitions.forEach(contentProviderDefinition -> {
                        if (validator.validate(processorManager, contentProviderDefinition)) {
                            WriterUtils.writeBaseDefinition(contentProviderDefinition, processorManager);
                        }
                    });
                }

                if(databaseHolderDefinition.databaseDefinition != null) {
                    databaseHolderDefinition.databaseDefinition.validateAndPrepareToWrite();
                }

                if (roundEnvironment.processingOver()) {
                    if(databaseHolderDefinition.databaseDefinition != null) {
                        if (databaseHolderDefinition.databaseDefinition.outputClassName != null) {
                            JavaFile.builder(databaseHolderDefinition.databaseDefinition.packageName, databaseHolderDefinition.databaseDefinition.typeSpec()).build()
                                    .writeTo(processorManager.processingEnvironment.getFiler());
                        }
                    }
                }

                List<TableDefinition> tableDefinitions = databaseHolderDefinition.tableDefinitionMap.values()
                    .stream().sorted(Comparator.comparing(tableDefinition -> tableDefinition.outputClassName.simpleName())).collect(Collectors.toList());

                tableDefinitions.forEach(tableDefinition -> WriterUtils.writeBaseDefinition(tableDefinition, processorManager));

                Collection<ModelViewDefinition> modelViewDefinitions = databaseHolderDefinition.modelViewDefinitionMap.values();
                modelViewDefinitions
                    .stream().sorted((a, b) -> a.priority > b.priority ? -1 : 1)
                    .forEach(modelViewDefinition -> WriterUtils.writeBaseDefinition(modelViewDefinition, processorManager));

                List<QueryModelDefinition> queryModelDefinitions = databaseHolderDefinition.queryModelDefinitionMap.values()
                    .stream().sorted(Comparator.comparing(queryModelDefinition -> queryModelDefinition.outputClassName.simpleName())).collect(Collectors.toList());
                queryModelDefinitions.forEach(queryModelDefinition -> WriterUtils.writeBaseDefinition(queryModelDefinition, processorManager) );

                EntityDefinition.safeWritePackageHelper(tableDefinitions, processorManager);
                EntityDefinition.safeWritePackageHelper(modelViewDefinitions, processorManager);
                EntityDefinition.safeWritePackageHelper(queryModelDefinitions, processorManager);
            } catch (IOException e) {
            }
        }

        try {
            DatabaseHolderDefinition databaseHolderDefinition = new DatabaseHolderDefinition(processorManager);
            if (!databaseHolderDefinition.isGarbage()) {
                JavaFile.builder(ClassNames.FLOW_MANAGER_PACKAGE,
                    databaseHolderDefinition.typeSpec()).build()
                    .writeTo(processorManager.processingEnvironment.getFiler());
            }
        } catch (FilerException ignored) {
        } catch (IOException e) {
            logError(e.getMessage());
        }
    }

    public boolean elementBelongsInTable(Element element) {
        Element enclosingElement = element.getEnclosingElement();
        EntityDefinition find = null;

        Collection<DatabaseObjectHolder> holders = databaseDefinitionMap.values();
        for (DatabaseObjectHolder holder : holders) {
            // table check.
            Collection<TableDefinition> tableDefinitions = holder.tableDefinitionMap.values();
            for(TableDefinition tableDefinition : tableDefinitions) {
                if(tableDefinition.element == enclosingElement) {
                    find = tableDefinition;
                    break;
                }
            }

            // modelview check.
            Collection<ModelViewDefinition> modelViewDefinitions = holder.modelViewDefinitionMap.values();
            for(ModelViewDefinition modelViewDefinition : modelViewDefinitions) {
                if(modelViewDefinition.element == enclosingElement) {
                    find = modelViewDefinition;
                    break;
                }
            }

            // querymodel check.
            Collection<QueryModelDefinition> queryModelDefinitions = holder.queryModelDefinitionMap.values();
            for(QueryModelDefinition queryModelDefinition : queryModelDefinitions) {
                if(queryModelDefinition.element == enclosingElement) {
                    find = queryModelDefinition;
                    break;
                }
            }
        }
        return find != null;
    }

}
